# -*- coding: utf-8 -*-

import configparser
import logging
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from Cal_MACD import calMACD


config = configparser.ConfigParser()
config_file = 'config.ini'
if not os.path.exists(config_file):
    raise FileNotFoundError(f"Configuration file '{config_file}' not found.")


config.read(config_file)
DATA_DIR = config['DEFAULT'].get('DataDir', 'd:/Python/study_data')
LIST_DIR = config['DEFAULT'].get('List', 'list')
LOG_FILE = config['DEFAULT'].get('LogFile', 'log.txt')

if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)

filelog = True
logger = logging.getLogger('log')
logger.setLevel(logging.DEBUG)
while logger.hasHandlers():
    for i in logger.handlers:
        logger.removeHandler(i)
# file log
# formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
formatter = logging.Formatter('%(message)s')

if filelog:
    fh = logging.FileHandler(LOG_FILE, encoding='utf-8')
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(formatter)
    logger.addHandler(fh)


def test_all_stocks(folder):
    hold_5_days = []
    hold_10_days = []
    hold_20_days = []
    hold_30_days = []
    hold_60_days = []
    for file in os.listdir(folder):
        if file.endswith('.csv'):
            path = os.path.join(folder, file)
            result = test_stock(path)
            hold_5_days.append(result[0])
            hold_10_days.append(result[1])
            hold_20_days.append(result[2])
            hold_30_days.append(result[3])
            hold_60_days.append(result[4])

    print(f"Average return for holding 5 days: {np.mean(hold_5_days):.4f}")
    print(f"Average return for holding 10 days: {np.mean(hold_10_days):.4f}")
    print(f"Average return for holding 20 days: {np.mean(hold_20_days):.4f}")
    print(f"Average return for holding 30 days: {np.mean(hold_30_days):.4f}")
    print(f"Average return for holding 60 days: {np.mean(hold_60_days):.4f}")


def test_stock(path):
    # 获取股票数据
    data = pd.read_csv(path)

    # 确保索引是日期时间类型
    if 'date' in data.columns:
        data.index = pd.to_datetime(data.date)
    else:
        raise ValueError("The 'date' column is missing from the data.")

    # 确保索引是日期时间类型
    data.index = pd.to_datetime(data.date)
    data = calMACD(data, 12, 26, 9)
    data.index = pd.to_datetime(data.date)
    data['20_day_min'] = data['close'].rolling(window=20).min()
    data['80_day_max'] = data['close'].rolling(window=80).max()

    data['buy_signal'] =  (data['dea9'] > -1.25) & (data['diff'] > -1.2) & (data['macd'].shift(1) < data['macd']) & (data['macd'] > 0) & (data['20_day_min'] < data['80_day_max'] * 0.7)
    data['buy_signal_'] = (data['buy_signal'] & ~(data['buy_signal'].shift(1) | data['buy_signal'].shift(2) | data['buy_signal'].shift(3) | data['buy_signal'].shift(4) | data['buy_signal'].shift(5)))
    # 检查生成的买入信号日期
    buy_signal_dates = data[data['buy_signal_']].index
    # print(f"Buy signal dates: {buy_signal_dates}")
    dict = {}
    dict[os.path.basename(path).split('.')[0]] = [str(date)[:10] for date in buy_signal_dates]

    # 初始化用于记录不同持有期收益的DataFrame
    hold_periods = [5, 10, 20, 30, 60]
    hold_days = [f'hold_{period}_days' for period in hold_periods]
    returns_df = pd.DataFrame(index=data.index, columns=hold_days)

    # 计算不同持有期的收益
    for period in hold_periods:
        for buy_date in buy_signal_dates:
            sell_date_idx = data.index.get_loc(buy_date) + period
            if sell_date_idx < len(data.index):
                buy_price = data.loc[buy_date, 'close']
                sell_date = data.index[sell_date_idx]
                sell_price = data.loc[sell_date, 'close']
                if pd.isna(sell_price):
                    logging.warning(f"Sell price is NaN for buy date {buy_date} and sell date {sell_date}")
                    returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan
                else:
                    returns_df.loc[buy_date, f'hold_{period}_days'] = (sell_price - buy_price) / buy_price
            else:
                # logging.warning(f"Sell date index {sell_date_idx} is out of range for buy date {buy_date}")
                returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan

    # 绘制不同持有期的累计收益曲线
    results = []
    for period in hold_periods:
        cumulative_returns = (1 + returns_df[f'hold_{period}_days']).cumprod(skipna=True) - 1
        plt.plot(cumulative_returns, label=f'Hold {period} days')
        avg_return = returns_df[f'hold_{period}_days'].mean(skipna=True)
        results.append(avg_return)
        # print(f"Average return for holding {period} days: {avg_return:.4f}")

    # print(dict)
    return dict


def find_all_buys(path):
    all_buys = {}
    all_buys_name = {}
    list_df = pd.read_csv(f'{LIST_DIR}/stocks_macd.csv')
    stocks = dict(zip(list_df['名称'], list_df['代码']))
    for name, code in stocks.items():
        if 'ST' not in name:
            try:
                buys = test_stock(f'{DATA_DIR}/{path}/{code}.csv')
                for key, value in buys.items():
                    for date in value:
                        if date in all_buys:
                            all_buys[date].append(key)
                            all_buys_name[date].append(name)
                        else:
                            all_buys[date] = [key]
                            all_buys_name[date] = [name]
            except FileNotFoundError:
                # print(f"File not found for stock {name} with code {code}")
                continue

    all_buys_df = pd.DataFrame.from_dict(all_buys, orient='index').sort_index()
    all_buys_name_df = pd.DataFrame.from_dict(all_buys_name, orient='index').sort_index()
    all_buys_df.to_csv('all_buys_macd.csv')
    all_buys_name_df.to_csv('all_buys_macd_name.csv')


if __name__ == '__main__':
    find_all_buys('stocks')
